import numpy as np
from matplotlib import pyplot as plt
from PIL import Image
import torch
import cv2
from utils.label_transform import rt2noised_rt

# plt显示中文
plt.rcParams["font.sans-serif"] = ["SimHei"]  # 设置字体
plt.rcParams["axes.unicode_minus"] = False  # 该语句解决图像中的“-”负号的乱码问题


# 线性变换回0-255
def norm_image(image, min_v=None, max_v=None, bits=None):
    if min_v is None:
        min_v = np.min(image)
    if max_v is None:
        max_v = np.max(image)
    if bits is None:
        return 255 * (image - min_v) / (max_v - min_v)
    else:
        return (2 ** bits - 1) * (image - min_v) / (max_v - min_v)


def norm_added_image(img, L=800):
    # added_img的大小为(4，x, y)
    _a = np.expand_dims(norm_image(img[0]) / 255, axis=0)
    _b = np.expand_dims(norm_image(img[1]) / 255, axis=0)
    _c = np.expand_dims(img[2] / L, axis=0)
    _d = np.expand_dims(norm_image(img[3]) / 255, axis=0)
    m1 = np.vstack((_a, _b))
    m2 = np.vstack((m1, _c))
    m3 = np.vstack((m2, _d))
    return m3


def normalize(v):
    mean = np.mean(v)
    std = np.std(v)
    if std == 0:
        v = v - mean
        return v
    return (v - mean) / std


def get_DRR_image(projector, tru_img, tru_r, tru_t, pre_r, pre_t, i, saving_path, mode):
    """
    获取实际的DRR和误差DRR之间的对比
    :param projector: 生成DRR的投影仪
    :param tru_img:实际图片
    :param tru_r: 实际的角度
    :param tru_t: 实际的移动
    :param pre_r: 预测的角度
    :param pre_t: 预测的移动
    :param i:测试的序号
    :param saving_path: 生成对比图的保存路径
    :param mode: 标准正位/标准侧位
    :return:
    """
    # 获取输入的DRR图像
    try:
        test_img = tru_img.detach().cpu().numpy().reshape([128, 128])
    except AttributeError:
        test_img = tru_img.reshape([128, 128])
    test_img = norm_image(test_img)
    # test_img = test_img*255
    truth_img = Image.fromarray(np.uint8(test_img))
    # 获取预测的DRR图像
    # 设置旋转角
    if mode == "标准正位":
        projector.set_rotation(0, 270, 90)
    elif mode == "标准侧位":
        projector.set_rotation(0, 90, 180)
    # 设置位移距离
    im_numpy = projector.project(angle_noise=rt2noised_rt(pre_r, mode), trans_noise=pre_t)
    im_numpy = normalize(im_numpy).reshape([128, 128])
    im_numpy = norm_image(im_numpy)
    pre_img = Image.fromarray(np.uint8(im_numpy))
    # 保存图片
    plt.figure(figsize=(10, 5))
    # 真实值
    plt.subplot(1, 2, 1)
    plt.title("tru-r{0}-t{1}".
              format([round(tru_r[0], 1), round(tru_r[1], 1), round(tru_r[2], 1)],
                     [round(tru_t[0], 1), round(tru_t[1], 1), round(tru_t[2], 1)]))
    plt.imshow(truth_img, cmap='gray')
    # 预测值
    plt.subplot(1, 2, 2)
    plt.title("pre-r{0}-t{1}".
              format([round(pre_r[0], 1), round(pre_r[1], 1), round(pre_r[2], 1)],
                     [round(pre_t[0], 1), round(pre_t[1], 1), round(pre_t[2], 1)]))
    plt.imshow(pre_img, cmap='gray')
    plt.suptitle("第{0}次测试".format(i))
    # print(test_img)
    # print(pre_img)
    # plt.show()
    if saving_path is not None:
        plt.savefig(saving_path + "/第{0}组.png".format(i))
    plt.close()


def get_DRR_image_re(projector, tru_img, pre_img, tru_r, tru_t, pre_r, pre_t, refine_r, refine_t, i, saving_path):
    """
    获取实际的DRR和误差DRR之间的对比以及第一次矫正之后的DRR
    :param projector: 生成DRR的投影仪
    :param tru_img:实际图片
    :param pre_img:第一次预测参数对应的DRR
    :param tru_r: 实际的角度
    :param tru_t: 实际的移动
    :param pre_r: 预测的角度
    :param pre_t: 预测的移动
    :param refine_r: 修正后的角度
    :param refine_t: 修正后的移动
    :param i:测试的序号
    :param saving_path: 生成对比图的保存路径
    :return:
    """
    # 获取输入的DRR图像
    test_img = tru_img.detach().cpu().numpy().reshape([128, 128])
    test_img = norm_image(test_img)
    truth_img = Image.fromarray(np.uint8(test_img))
    # 获取预测的DRR图像
    pre_img = pre_img.detach().cpu().numpy().reshape([128, 128])
    pre_img = norm_image(pre_img)
    pre_img = Image.fromarray(np.uint8(pre_img))
    # 设置旋转角
    projector.set_rotation(refine_r[0], refine_r[1], refine_r[2])
    # 设置位移距离
    im_numpy = projector.project(refine_t[0], refine_t[1], refine_t[2])
    im_numpy = normalize(im_numpy).reshape([128, 128])
    im_numpy = norm_image(im_numpy)
    refine_img = Image.fromarray(np.uint8(im_numpy))
    # 保存图片
    plt.figure(figsize=(15, 5))
    # 真实值
    plt.subplot(1, 3, 1)
    plt.title("实际-旋转{0}-平移{1}".format(np.int8(tru_r), np.int8(tru_t)))
    plt.imshow(truth_img, cmap='gray')
    # 预测值
    plt.subplot(1, 3, 2)
    plt.title("预测-旋转{0}-平移{1}".format(np.int8(pre_r), np.int8(pre_t)))
    plt.imshow(pre_img, cmap='gray')
    # 预测值
    plt.subplot(1, 3, 3)
    plt.title("修正-旋转{0}-平移{1}".format(np.int8(refine_r), np.int8(refine_t)))
    plt.imshow(refine_img, cmap='gray')
    plt.suptitle("第{0}次测试".format(i))
    plt.savefig(saving_path + "/第{0}组.png".format(i))
    plt.close()


def get_new_drr(rx, ry, rz, tx, ty, tz, size=8, data_generator=None):
    all_img = []
    for k in range(size):
        # 设置旋转角
        data_generator.set_rotation(rx[k, 0], ry[k, 0], rz[k, 0])
        # 设置位移距离
        im_numpy = data_generator.project(tx[k, 0], ty[k, 0], tz[k, 0])
        all_img.append(im_numpy / 255)
    all_img = np.array(all_img).reshape([size, 1, 128, 128])
    img_tensor = torch.from_numpy(all_img).float().cuda()
    return img_tensor


def get_new_drr_numpy(rx, ry, rz, tx, ty, tz, mode=None, data_generator=None):
    # 设置旋转角
    if mode == "标准正位":
        data_generator.set_rotation(0, 270, 90)
    elif mode == "标准侧位":
        data_generator.set_rotation(0, 90, 180)
    im_numpy = data_generator.project(angle_noise=np.array([rx, ry, rz]), trans_noise=np.array([tx, ty, tz]))
    return im_numpy


def img_translation(img, tx, ty):
    rows, cols, _ = img.shape
    # 平移矩阵M：[[1,0,x],[0,1,y]]
    M = np.float32([[1, 0, tx], [0, 1, ty]])
    dst = cv2.warpAffine(img, M, (cols, rows))
    return dst


def caculate_NCC(image1, image2):
    return np.mean(np.multiply((image1 - np.mean(image1)), (image2 - np.mean(image2)))) \
        / (np.std(image1) * np.std(image2))


def caculate_MSE(image1, image2):
    return np.mean((image1 - image2) ** 2)


def add_o_source_coord(img, A=None, L=800):
    _, u, v = img.shape
    coordinate = np.zeros((u * v, 3))
    if A is None:
        A = np.array([[-568.8889, 0, 63.5],
                      [0, 568.8889, 63.5],
                      [0, 0, 1.]]
                     )
    A_inv = np.linalg.inv(A)
    _id = 0
    for i in range(u):
        for j in range(v):
            uv = np.array([[L * i],
                           [L * j],
                           [L]]).astype(np.float64)
            coord = np.dot(A_inv, uv)
            # print(coordinate[_id])
            coordinate[_id] = coord.ravel()
            _id += 1
    coordinate = coordinate.reshape((u, v, 3)).transpose(2, 0, 1)
    added_img = np.vstack((img, coordinate))
    return added_img


if __name__ == "__main__":
    coord = np.load("../data/coord/normed_coord_xy.npy")
    # coord = coord[:2].reshape(2, 128, 128)
    # coord = normalize(coord)/255
    # np.save("C:/Users/adminTKJ/Desktop/RLIR_sumup/data/coord/normed_coord_xy.npy", coord)
    print(coord)
